#ifndef PCPS_ACCUMULATOR_ENERGY_HEADER
#define PCPS_ACCUMULATOR_ENERGY_HEADER

#include <src/pcps.h>
#include <src/input.h>
#include <src/hamiltonian.h>
#include <src/wavefunction.h>
#include <src/accumulator.h>
#include <src/local_energy.h>

namespace pcps {

  // forward declaration of thread accumulator
  template <class S, class PART, class REF> class EnergyThreadAccum;

  //--------------------------------------------------------------------------------
  // pcps::EnergyProcessAccum -- Accumulates energy data for one process.
  //--------------------------------------------------------------------------------

  template <class S, class PART, class REF> class EnergyProcessAccum : public pcps::ProcessAccum<S, PART, REF> {

    private:

      // data members
      const int _nsites;
      S _total_energy;
      S _accept_ratio;
      const pcps::LocalEnergyFunc<S, PART, REF> _energy_func;

    public:

      // give the corresponding thread accumulator access to private and protected members
      friend class pcps::EnergyThreadAccum<S, PART, REF>;

      // constructor
      EnergyProcessAccum(const pcps::Input & userinp, const pcps::Hamiltonian<S> & ham, const pcps::Wavefunction<S> & wfn, const long int sl)
        : pcps::ProcessAccum<S, PART, REF>(userinp, sl),
          _nsites(userinp.nsites()),
          _energy_func(userinp, ham, wfn)
      {
        this->reset();
      }

      // function to create a thread accumulator linked to this process accumulator
      pcps::ThreadAccum<S, PART, REF> * create_thread_accum(const int tid, pcps::Arena & arena);

      // function to reset the accumulator
      void reset() {
        _total_energy = pcps::zero<S>();
        _accept_ratio = pcps::zero<S>();
      }

      // function to return the acceptance ratio
      double acceptance_ratio() const { return pcps::real(_accept_ratio); }

      // function to return the total energy from the process accumulator
      S total_energy() const { return _total_energy; }

      // function to print information during stochastic iterations
      void print_iter_info(const pcps::Input & userinp,
                           pcps::Wavefunction<S> & wfn,
                           const int iter,
                           std::vector<boost::shared_ptr<pcps::ProcessAccum<S, PART, REF> > > & accumulators) {

        if (this->myrank() == 0) {
          std::cout << boost::format("energy                 %20.12f") % pcps::real(this->total_energy()) << std::endl;
          std::cout << boost::format("acceptance ratio       %20.12f") % this->acceptance_ratio() << std::endl;
        }

      }

  };

  //--------------------------------------------------------------------------------
  // pcps::EnergyThreadAccum -- Accumulates energy data for one thread.
  //--------------------------------------------------------------------------------

  template <class S, class PART, class REF> class EnergyThreadAccum : public pcps::ThreadAccum<S, PART, REF> {

    private:

      // pointer to the corresponding process accumulator
      pcps::EnergyProcessAccum<S, PART, REF> * const _epacc;

      // data members
      S _local_energy;
      S _energy;
      long int _n_accept;
      int * _workspace;
      S * _reduce_space;

    public:

      // constructor
      EnergyThreadAccum(pcps::EnergyProcessAccum<S, PART, REF> * const epacc, const int tid, pcps::Arena & arena)
        : pcps::ThreadAccum<S, PART, REF>(epacc, tid),
          _epacc(epacc)
      {
        _workspace = arena.allocate_array<int>( 2 * size_t(_epacc->_nsites), 0);
        if (this->_tid == 0)
          _reduce_space = arena.allocate_array<S>(4, pcps::zero<S>());
        this->reset();
      }

      // function to reset the accumulator
      void reset() {
        _local_energy = pcps::zero<S>();
        _energy = pcps::zero<S>();
        _n_accept = 0;
      }

      // function to return the local energy
      S local_energy() const { return _local_energy; }

      // function to return the total energy from the process accumulator
      S total_energy() const { return _epacc->_total_energy; }

      // function to accumulate the energy of the old configuration
      void accumulate(const long int iter, 
                      const long int count,
                      const bool accept,
                      const bool finished,
                      const S coeff_ratio,
                      const int * const new_config,
                      const int * const old_config,
                      pcps::CoeffValueFunc<S, PART, REF> & cvf,
                      pcps::ThreadAccum<S, PART, REF> ** accumulators,
                      const int n_accumulators) {

        if (accept)
          _n_accept++;

        if (accept || finished) {
          _local_energy = _epacc->_energy_func.get_local_energy(old_config, _workspace, cvf);
          _energy += double(count) * _local_energy;
        }

        if (finished) {

          // add the thread accumulator's energy and number of accepted moves to the process accumulator
          pthread_mutex_lock(this->_pacc->mutex());
          _epacc->_total_energy += _energy;
          _epacc->_accept_ratio += pcps::unity<S>() * double(_n_accept);
          pthread_mutex_unlock(this->_pacc->mutex());

          // wait for other threads to add their energies
          pthread_barrier_wait(this->_pacc->barrier());

          // thread 0 handles communication between processes
          if (this->_tid == 0) {

            // reduce energy and acceptance ratio on all processes
            _reduce_space[0] = _epacc->_total_energy;
            _reduce_space[1] = _epacc->_accept_ratio;
            _reduce_space[2] = pcps::zero<S>();
            _reduce_space[3] = pcps::zero<S>();
            pcps::allreduce(MPI::COMM_WORLD, _reduce_space, _reduce_space+2, 2, MPI::SUM);
            _epacc->_total_energy = _reduce_space[2] / this->_pacc->total_samples();
            _epacc->_accept_ratio = _reduce_space[3] / this->_pacc->total_samples();

          }

          // wait for thread 0 to finish communication
          pthread_barrier_wait(this->_pacc->barrier());

        }

      }

  };

  //------------------------------------------------------------------------------------------
  // pcps::EnergyProcessAccum::create_thread_accum -- Function to create a thread accumulator
  //                                                  from a process accumulator.
  //------------------------------------------------------------------------------------------

  template <class S, class PART, class REF>
  pcps::ThreadAccum<S, PART, REF> * pcps::EnergyProcessAccum<S, PART, REF>::create_thread_accum(const int tid, pcps::Arena & arena) {
    return new(arena) pcps::EnergyThreadAccum<S, PART, REF>(this, tid, arena);
  }

} // end namespace pcps

#endif
